# model.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models
import os

class CBAM(nn.Module):
    def __init__(self, in_channels, reduction=16):
        super().__init__()
        self.channel = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, in_channels // reduction, 1),
            nn.ReLU(),
            nn.Conv2d(in_channels // reduction, in_channels, 1),
            nn.Sigmoid()
        )
        self.spatial = nn.Sequential(
            nn.Conv2d(2, 1, kernel_size=7, padding=3),
            nn.Sigmoid()
        )

    def forward(self, x):
        c_weight = self.channel(x)
        x = x * c_weight
        avg = torch.mean(x, dim=1, keepdim=True)
        max = torch.max(x, dim=1, keepdim=True)[0]
        s_weight = self.spatial(torch.cat([avg, max], dim=1))
        return x * s_weight

class ResBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, stride),
            nn.Conv2d(out_channels, out_channels, 3, padding=1),
            nn.Conv2d(out_channels, out_channels, 1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
        self.downsample = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, stride),
            nn.BatchNorm2d(out_channels)
        ) if in_channels != out_channels else nn.Identity()

    def forward(self, x):
        return F.relu(self.block(x) + self.downsample(x))

class NLB(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.theta = nn.Conv2d(in_channels, in_channels//2, 1)
        self.phi = nn.Conv2d(in_channels, in_channels//2, 1)
        self.g = nn.Conv2d(in_channels, in_channels//2, 1)
        self.out = nn.Conv2d(in_channels//2, in_channels, 1)

    def forward(self, x):
        B, C, H, W = x.size()
        theta = self.theta(x).view(B, C//2, -1)
        phi = self.phi(x).view(B, C//2, -1)
        g = self.g(x).view(B, C//2, -1)
        attn = F.softmax(torch.bmm(theta.transpose(1,2), phi), dim=-1)
        out = torch.bmm(g, attn.transpose(1,2)).view(B, C//2, H, W)
        return x + self.out(out)

class DRNet(nn.Module):
    def __init__(self, num_classes=5):
        super().__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, 7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3, 2, 1)
        )
        self.res1 = nn.Sequential(*[ResBlock(64, 256) for _ in range(3)])
        self.cbam1 = CBAM(256)
        self.res2 = nn.Sequential(*[ResBlock(256, 512) for _ in range(4)])
        self.cbam2 = CBAM(512)
        self.res3 = nn.Sequential(*[ResBlock(512, 1024) for _ in range(6)])
        self.res4 = nn.Sequential(*[ResBlock(1024, 2048) for _ in range(3)])
        self.nlb = NLB(2048)
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(2048, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.res1(x)
        x = self.cbam1(x)
        x = self.res2(x)
        x = self.cbam2(x)
        x = self.res3(x)
        x = self.res4(x)
        x = self.nlb(x)
        x = self.pool(x)
        x = self.classifier(x)
        return x


def load_model():
    model_path = 'best.pth'
    num_classes = 5
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 创建模型结构
    model = models.resnet50(pretrained=False)
    model.fc = nn.Linear(model.fc.in_features, num_classes)
    model = model.to(device)
    model.eval()
    
    # 尝试加载权重
    try:
        if os.path.exists(model_path):
            checkpoint = torch.load(model_path, map_location=device)
            model.load_state_dict(checkpoint['model'])
            print(f"成功加载模型权重: {model_path}")
        else:
            print(f"模型文件不存在: {model_path}")
    except Exception as e:
        print(f"加载模型权重时出错: {e}")
    
    return model